package com.xncoding.websocket.sse.controller;

import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping("/sse")
public class SseController {

    // 该集合用来管理所有客户端的连接
    private final Map<String, SseEmitter> sse = new ConcurrentHashMap<>() ;

    // 创建连接接口，同时指定了消息类型为text/event-stream
    @GetMapping(path="/events/{id}", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter createConnect(@PathVariable("id") String id) throws IOException {
        SseEmitter emitter = new SseEmitter(0L);
        // 每一个客户端保存到Map中
        sse.put(id, emitter) ;
        // 当发生错误的回调
        emitter.onError(ex -> {
            System.err.printf("userId: %s, error: %s%n", id, ex.getMessage()) ;
            sse.remove(id) ;
        }) ;
        // 异步请求完成后的回调
        emitter.onCompletion(() -> {
            sse.remove(id) ;
            System.out.printf("%s, 请求完成...") ;
        }) ;
        // 异步请求超时回调
        emitter.onTimeout(() -> {
            System.err.println("超时...") ;
        }) ;
        return emitter;
    }

    // 该接口用来进行消息的发送
    // 由客户端发起请求，然后根据id获取相应的SseEmitter进行消息的发送
    @GetMapping("/sender/{id}")
    public String sender(@PathVariable("id") String id) throws Exception {
        SseEmitter emitter = this.sse.get(id) ;
        if (emitter != null) {
            SseEmitter.SseEventBuilder builder = SseEmitter.event() ;
            // 指定事件类型
            builder.name("chat") ;
            String msg = "随机消息 - " + new Random().nextInt(10000000);
            builder.data(msg) ;
            try {
                emitter.send(builder) ;
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return "success" ;
    }
}